from typing import Dict, Any, Tuple
import pandas as pd
import numpy as np
from .moving_average_strategy import MovingAverageStrategy
from ..risk.risk_manager import RiskManager
from ..config.config import Config

class RiskAwareMAStrategy(MovingAverageStrategy):
    def __init__(self, config: Config, parameters: Dict[str, Any] = None):
        super().__init__(parameters)
        self.risk_manager = RiskManager(config)
        self.config = config

    def generate_signals(self, data: pd.DataFrame) -> Tuple[pd.Series, Dict]:
        """
        Generate trading signals with risk management parameters
        """
        # Get basic signals from parent class
        basic_signals = super().generate_signals(data)
        
        # Calculate volatility (using ATR)
        atr = self._calculate_atr(data)
        
        # Initialize risk-adjusted signals and position sizes
        risk_adjusted_signals = pd.Series(0, index=data.index)
        position_sizes = pd.Series(0.0, index=data.index)
        stop_losses = pd.Series(0.0, index=data.index)
        
        for i in range(len(data)):
            if basic_signals.iloc[i] != 0:
                # Calculate position size based on risk
                position_size = self.risk_manager.calculate_position_size(
                    symbol=data.index[i].symbol if hasattr(data.index[i], 'symbol') else 'Unknown',
                    current_price=data['Close'].iloc[i],
                    volatility=atr.iloc[i],
                    available_cash=self.risk_manager.portfolio_value
                )
                
                # Calculate stop loss
                stop_loss = self.risk_manager.calculate_stop_loss(
                    entry_price=data['Close'].iloc[i],
                    position_type='long' if basic_signals.iloc[i] == 1 else 'short',
                    atr=atr.iloc[i]
                )
                
                # Only generate signal if position size is meaningful
                if position_size > 0:
                    risk_adjusted_signals.iloc[i] = basic_signals.iloc[i]
                    position_sizes.iloc[i] = position_size
                    stop_losses.iloc[i] = stop_loss
        
        return risk_adjusted_signals, {
            'position_sizes': position_sizes,
            'stop_losses': stop_losses,
            'atr': atr
        }

    def _calculate_atr(self, data: pd.DataFrame, period: int = 14) -> pd.Series:
        """
        Calculate Average True Range (ATR)
        """
        high = data['High']
        low = data['Low']
        close = data['Close']
        
        tr1 = high - low
        tr2 = abs(high - close.shift())
        tr3 = abs(low - close.shift())
        
        tr = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        atr = tr.rolling(window=period).mean()
        
        return atr 